// wsp 2 tcp 源码来自https://github.com/xjdrew/wsg
package main

import (
	"bufio"
	"bytes"
	"encoding/binary"
	"errors"
	"flag"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"runtime/debug"
	"sync"
	"sync/atomic"
	"time"

	"github.com/gorilla/websocket"
	log "github.com/sirupsen/logrus"
	"gopkg.in/natefinch/lumberjack.v2"
)

const (
	// Time allowed to write a message to the peer.
	writeWait = 10 * time.Second

	// Time allowed to read the next pong message from the peer.
	pongWait = 60 * time.Second

	// Send pings to peer with this period. Must be less than pongWait.
	pingPeriod = (pongWait * 9) / 10

	// Maximum message size allowed from peer.
	maxMessageSize = 4096
)

var (
	maxConnCount         = 2000 // 最大运行连接数
	connCount            = int32(0)
	websocketMessageType = websocket.TextMessage
	splitPackageLen      = 2 // 粘包、拆包长度 支持2个字节和4个字节
	upgrader             = websocket.Upgrader{
		ReadBufferSize:  4096,
		WriteBufferSize: 4096,
		CheckOrigin: func(r *http.Request) bool {
			return true
		},
	}
)

type Client struct {
	clientConn *websocket.Conn // The websocket connection from client
	serverConn *ServerConn     // connection to server
	closeOnce  sync.Once
	sendCount  int64
	recvCount  int64
}

// readPump pumps messages from the websocket connection to the server.
func (c *Client) readPump() {
	defer func() {
		log.Warnf("client remoteaddr:%s localaddr:%s quit client transfer to server coroutine ", c.clientConn.RemoteAddr().String(), c.serverConn.LocalAddr().String())
		c.close()
	}()
	c.clientConn.SetReadLimit(maxMessageSize)
	c.clientConn.SetReadDeadline(time.Now().Add(pongWait))
	c.clientConn.SetPongHandler(func(string) error { c.clientConn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
	for {
		_, message, err := c.clientConn.ReadMessage()
		if err != nil {
			log.Errorf("client remoteaddr:%s localaddr:%s in read from client error:%v", c.clientConn.RemoteAddr().String(), c.serverConn.LocalAddr().String(), err)
			if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
			}
			break
		}
		c.clientConn.SetReadDeadline(time.Now().Add(pongWait))
		log.Debugf("client remoteaddr:%s localaddr:%s read from client len:%d pack:%v", c.clientConn.RemoteAddr().String(), c.serverConn.LocalAddr().String(), len(message), message)
		if websocketMessageType == websocket.BinaryMessage {
			buf := new(bytes.Buffer)
			if splitPackageLen == 2 {
				err = binary.Write(buf, binary.LittleEndian, uint16(len(message))) // 添加2个字节填充长度
				if err != nil {
					log.Errorf("client remoteaddr:%s localaddr:%s in readPump write len error:%v", c.clientConn.RemoteAddr().String(), c.serverConn.LocalAddr().String(), err)
					break
				}
			} else if splitPackageLen == 4 {
				err = binary.Write(buf, binary.LittleEndian, uint32(len(message))) // 添加4个字节填充长度
				if err != nil {
					log.Errorf("client remoteaddr:%s localaddr:%s in readPump write len error:%v", c.clientConn.RemoteAddr().String(), c.serverConn.LocalAddr().String(), err)
					break
				}
			} else {
				log.Errorf("client remoteaddr:%s localaddr:%s in readPump write len error:invalid splitPackageLen", c.clientConn.RemoteAddr().String(), c.serverConn.LocalAddr().String())
				break
			}

			buf.Write(message)
			if _, err = c.serverConn.Write(buf.Bytes()); err != nil {
				log.Errorf("client remoteaddr:%s localaddr:%s in write to server error:%v", c.clientConn.RemoteAddr().String(), c.serverConn.LocalAddr().String(), err)
				break
			}
			c.sendCount++
		} else {
			if _, err = c.serverConn.Write(message); err != nil {
				log.Errorf("client remoteaddr:%s localaddr:%s in write to server error:%v", c.clientConn.RemoteAddr().String(), c.serverConn.LocalAddr().String(), err)
				break
			}
			c.sendCount++
		}
	}
}

// writePump pumps messages from the server to the websocket connection.
func (c *Client) writePump() {
	defer func() {
		log.Warnf("client remoteaddr:%s localaddr:%s quit server transfer to client coroutine ", c.clientConn.RemoteAddr().String(), c.serverConn.LocalAddr().String())
		c.close()
	}()

	input := bufio.NewScanner(c.serverConn)
	// NewScanner 默认以  进行拆包，如果是二进制则以以下定义的格式进行拆包(2个字节长度)
	if websocketMessageType == websocket.BinaryMessage {
		split := func(data []byte, atEOF bool) (advance int, token []byte, err error) {
			if !atEOF {
				if len(data) > splitPackageLen {
					var dataLen int16
					binary.Read(bytes.NewReader(data[0:splitPackageLen]), binary.LittleEndian, &dataLen)
					if dataLen < 0 {
						log.Errorf("client remoteaddr:%s localaddr:%s split read datalen:%d error ", c.clientConn.RemoteAddr().String(), c.serverConn.LocalAddr().String(), dataLen)
						return 0, []byte(""), errors.New("read datalen err")
					}
					if int(dataLen)+splitPackageLen <= len(data) {
						return int(dataLen) + splitPackageLen, data[splitPackageLen : int(dataLen)+splitPackageLen], nil
					}
				}
			}
			return
		}
		input.Split(split)
	}
	for input.Scan() { //会阻塞等待客户端数据到来
		c.clientConn.SetWriteDeadline(time.Now().Add(writeWait))
		log.Debugf("client remoteaddr:%s localaddr:%s read from server len:%d pack:%v", c.clientConn.RemoteAddr().String(), c.serverConn.LocalAddr().String(), len(input.Bytes()), input.Bytes())
		if err := c.clientConn.WriteMessage(websocketMessageType, input.Bytes()); err != nil {
			log.Errorf("client remoteaddr:%s localaddr:%s in write to client error:%v", c.clientConn.RemoteAddr().String(), c.serverConn.LocalAddr().String(), err)
			break
		}
		c.recvCount++
	}
}

func (c *Client) run() {
	go c.readPump()
	go c.writePump()
}

func (c *Client) close() {
	c.closeOnce.Do(func() {
		atomic.AddInt32(&connCount, -1)
		log.Infof("client offile remoteaddr %s localaddr:%s sendcount:%d recvcount:%d conncount:%d", c.clientConn.RemoteAddr().String(), c.serverConn.LocalAddr().String(), c.sendCount, c.recvCount, atomic.LoadInt32(&connCount))
		c.serverConn.Close()
		c.clientConn.Close()
	})
}

// serveWs handles websocket requests from the peer.
func serveWs(server *Server, w http.ResponseWriter, r *http.Request) {
	clientConn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	serverConn, err := server.Dial()
	if err != nil {
		clientConn.Close()
		log.Println(err)
		return
	}

	atomic.AddInt32(&connCount, 1)
	log.Infof("client online remoteaddr:%s localaddr:%s conncount:%d", clientConn.RemoteAddr().String(), serverConn.LocalAddr().String(), atomic.LoadInt32(&connCount))
	client := &Client{clientConn: clientConn, serverConn: serverConn}
	client.run()
}

type ServerConn struct {
	net.Conn
}

type Server struct {
	addr string
}

func (server Server) Dial() (*ServerConn, error) {
	conn, err := net.Dial("tcp", server.addr)
	if err != nil {
		return nil, err
	}
	return &ServerConn{conn}, nil
}

func newServer(addr string) *Server {
	return &Server{addr: addr}
}

func dump(id string) {
	rcv_err := recover()
	if rcv_err == nil {
		return
	}
	now := time.Now()
	strTime := now.Format("20060102150405")
	name := fmt.Sprintf("./log/dump-%s-%s.log", id, strTime)
	fmt.Println("dump file in", name)
	f, err := os.Create(name)
	if err != nil {
		return
	}
	defer f.Close()
	f.WriteString(fmt.Sprintf("panic:\r\n%v", rcv_err))
	f.WriteString("\r\n\nstack:\r\n")
	f.WriteString(string(debug.Stack()))
}

func main() {
	// 初始化命令行参数
	var (
		id         string
		wsAddr     string
		tcpAddr    string
		binaryMode bool

		logPath  string
		logLevel string
	)

	flag.StringVar(&id, "id", fmt.Sprintf("%d", os.Getpid()), "tag for exe")
	flag.StringVar(&wsAddr, "ws", ":4223", "websocket server address")
	flag.StringVar(&tcpAddr, "tcp", "127.0.0.1:4224", "backend tcp server address")
	flag.IntVar(&splitPackageLen, "split", 2, "split package len")
	flag.BoolVar(&binaryMode, "binary", false, "Use binary frames instead of text frames")
	flag.StringVar(&logPath, "logpath", "./log/", "log path")
	flag.StringVar(&logLevel, "loglevel", "info", "log level")
	flag.Parse()

	if binaryMode {
		websocketMessageType = websocket.BinaryMessage
	}

	defer dump(id)

	// 初始化日志
	fileName := fmt.Sprintf("%s%s.log", logPath, id)
	logger := &lumberjack.Logger{
		// 日志输出文件路径
		Filename: fileName,
		// 日志文件最大 size, 单位是 MB
		MaxSize: 100, // megabytes
		// 最大过期日志保留的个数
		MaxBackups: 30,
		// 保留过期文件的最大时间间隔,单位是天
		MaxAge: 7, //days
		// 是否需要压缩滚动日志, 使用的 gzip 压缩
		Compress: true, // disabled by default
	}
	writers := []io.Writer{logger, os.Stdout}
	mutliWriter := io.MultiWriter(writers...)
	log.SetOutput(mutliWriter)
	if l, err := log.ParseLevel(logLevel); err == nil {
		fmt.Printf("set log level:%d\n", l)
		log.SetLevel(l)
	}
	log.Infof("log init ...............................................")

	// 启动websocket服务器
	log.Infof("Listen on ws addr:%s and backend tcp addr:%s binary:%t loglevel:%s", wsAddr, tcpAddr, binaryMode, logLevel)
	server := newServer(tcpAddr)
	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		serveWs(server, w, r)
	})
	if err := http.ListenAndServe(wsAddr, nil); err != nil {
		log.Fatalf("ListenAndServe: %v", err)
	}
}
